import requests
import pandas as pd
import time
import json


def fetch_sz_stock_list(page, page_size):
    url = "https://push2.eastmoney.com/api/qt/clist/get"
    params = {
        'np': 1,
        'fltt': 1,
        'invt': 2,
        'cb': 'jQuery37103545651209258376_1743989084102',  # 可以保留，或者每次生成一个新的
        'fs': 'm:0+t:6,m:0+t:80',  # 深证A股
        'fields': 'f12,f13,f14,f1,f2,f4,f3,f152,f5,f6,f7,f15,f18,f16,f17,f10,f8,f9,f23',
        'fid': 'f12',
        'pn': page,  # 页码
        'pz': page_size,  # 每页数量
        'po': 0,
        'dect': 1,
        'ut': 'fa5fd1943c7b386f172d6893dbfba10b',
        'wbp2u': '|0|0|0|web',
        '_': str(int(time.time() * 1000))  # 当前时间戳
    }

    response = requests.get(url, params=params)

    if response.status_code == 200:
        # 去掉回调函数的部分
        json_str = response.text[response.text.index('(') + 1:-2]
        try:
            data = json.loads(json_str)  # 解析 JSON 数据
            return data  # 返回解析后的数据
        except ValueError as e:
            print("JSON解析错误:", e)
            return None
    else:
        print("请求失败，状态码：", response.status_code)
        return None


def save_to_csv(stock_data, filename):
    # 将数据保存到 CSV 文件
    df = pd.DataFrame(stock_data)
    df.to_csv(filename, index=False, encoding='utf-8-sig')
    print(f"数据已保存到 {filename}")


def main():
    all_stocks = []
    page_size = 20  # 每页数量
    page = 1

    # 首先获取第一页数据以获取总股票数量
    print("正在获取第 1 页数据以获取总股票数量...")
    first_page_data = fetch_sz_stock_list(page, page_size)

    if first_page_data and 'data' in first_page_data:
        total_stocks = first_page_data['data']['total']  # 获取总股票数量
        total_pages = (total_stocks + page_size - 1) // page_size  # 计算总页数

        print(f"总股票数量: {total_stocks}, 总页数: {total_pages}")

        # 从第一页开始遍历所有股票数据
        for page in range(1, total_pages + 1):
            print(f"正在获取第 {page} 页数据...")
            stock_data = fetch_sz_stock_list(page, page_size)

            if stock_data and 'data' in stock_data:
                stocks = stock_data['data']['diff']  # 提取股票数据
                all_stocks.extend(stocks)
                print(f"成功获取第 {page} 页数据，当前股票总数：{len(all_stocks)}")
                time.sleep(1)  # 为了避免请求过于频繁，暂停 1 秒
            else:
                print(f"第 {page} 页数据获取失败，尝试下一页。")

        # 保存获取的所有股票数据到 CSV 文件
        if all_stocks:
            save_to_csv(all_stocks, 'data/sz_ashares.csv')  # 修改文件名
        else:
            print("没有获取到任何股票数据。")
    else:
        print("获取第一页数据失败，无法获得总股票数量。")


if __name__ == '__main__':
    main()